from openai import AsyncOpenAI
import asyncio
from typing_extensions import override
from .BasePool import BasePool, Entity, PoolConfig


class OpenAiPool(BasePool[AsyncOpenAI]):
    def __init__(self, base_url: str, api_key: str, config: PoolConfig = None):
        super().__init__(config)
        self.base_url = base_url
        self.api_key = api_key

    @override
    async def create_connection(self):
        if len(self._pool) >= self.config.max_connections:
            self.config.logger.warning(
                f"已达到最大连接数 ({self.config.max_connections}), 无法创建新连接"
            )
            return None
        for attempt in range(self.config.connection_retry + 1):
            try:
                # 创建连接
                client = AsyncOpenAI(
                    base_url=self.base_url,
                    api_key=self.api_key,
                )

                connection = Entity[AsyncOpenAI](client=client)

                self.config.logger.debug(
                    f"成功创建新的ChatSession连接 (当前连接数: {len(self._pool) + 1})"
                )
                if await self.validate_connection(connection):
                    return connection

            except Exception as e:
                if attempt < self.config.connection_retry:
                    self.config.logger.warning(
                        f"创建ChatSession失败 (尝试 {attempt + 1}/{self.config.connection_retry + 1}): {str(e)}"
                    )
                    await asyncio.sleep(self.config.retry_interval)
                else:
                    self.config.logger.error(
                        f"创建ChatSession失败，已达到最大重试次数: {str(e)}"
                    )
                    return None

    async def validate_connection(self, connection: Entity[AsyncOpenAI]) -> bool:
        try:
            # 验证连接
            await connection.client.models.list()
            return True
        except Exception as e:
            self.config.logger.warning(f"验证ChatSession失败: {str(e)}")
            return False
